functions {
    real sqsum (vector x){
    real y = 0;
    
    for (i in 1:num_elements(x)){
      y = y + x[i]*x[i];
    }
    return y;
   }
  
  real pow_sum (vector x, real z){
    real y = 0;
    
    for (i in 1:num_elements(x)){
      y = y + x[i]^z;
    }
    return y;
   } 
  
   real calc_cor (vector x, vector y){
     real r;
     r = sum((x - mean(x)).*(y - mean(y))) / sqrt(sqsum(x - mean(x))*sqsum(y - mean(y)));
     return r;
   }
}
data {
  int<lower=1> I;               // proposals
  int<lower=1> J;               // legislators
  int<lower=1> N;               // nr of observed rc
  int<lower=1> J_old;           // legislators
  int<lower=1> D;               // number of latent dimensions
  int<lower=1> K;               // number of exp. variables
  int<lower=0,upper=1> T;       // include time component
  int<lower=-1,upper=1> Y[J,I]; // binary y
  int<lower=1,upper=I> ii[N];
  int<lower=1,upper=J> jj[N];
  real theta_in[J,K,D];         // explanatory variables for ideal points
  int<lower=0,upper=1> incl_nonsep;  // indicator to use separable or nonseparable model
  int<lower=1,upper=J> theta_old_ind[J_old];
  real theta_old[J_old,D];
}
parameters {
  simplex[D] sal; // salience of dimensions
  real<lower=0> tau; // scale of weights matrix overall
  real<lower=0> tau_int; // scale of salience internally in nonseparable setting
  matrix[J,D] theta_free; // free ideal points
  matrix[I,D] bk_free; // free discrimination
  vector[I] ak; // difficulty
  cholesky_factor_corr[D] L; // correlation component of weight matrix
  vector<lower=0>[D] gamma; // coefficients for linear regression on ideal points, lower bound to prevent switching
}
transformed parameters {
  corr_matrix[D] DL; // correlation component of weight matrix
  cov_matrix[D] weights; // final weight matrix
  matrix[J,D] theta_t;
  matrix[I,D] bk; // standardized discrimination (mean 0, var 1)
  matrix[J,D] theta; // standardized ideal points (var 1)

  DL = L * L';
  
  {
    vector[D] sd_theta;
    vector[D] sd_bk;
    vector[D] mean_theta;
    
    // it needs this constraint of sum(abs(...)) == 1 here so that weights is identified internally.
    // otherwise, same variance can be gotten by varying either var or cov within the final matrix
    if (incl_nonsep == 1){
      weights = quad_form_diag(DL, sqrt(sal * tau_int)) / sum(fabs(to_vector(quad_form_diag(DL, sqrt(sal * tau_int)))))*tau^2;
    } else {
      weights = quad_form_diag(diag_matrix(rep_vector(1, D)), sqrt(sal * tau));
    }
    
    for (d in 1:D){
      theta_t[,d] = to_vector(theta_in[,1,d]) * (gamma[d]) + to_vector(theta_in[,2,d]) + theta_free[,d];
    }
    
    for (d in 1:D){
      sd_theta[d] = sd(theta_t[,d]);
      sd_bk[d] = sd(bk_free[,d]);
      mean_theta[d] = mean(theta_t[,d]);
    }

    for (d in 1:D){
      theta[,d] = (theta_t[,d] - mean_theta[d]) / sd_theta[d];
      bk[,d] = bk_free[,d] / sd_bk[d];
    }
  }
}
model {
  to_vector(bk_free) ~ std_normal();
  ak ~ std_normal();
  sal ~ dirichlet(rep_vector(1.00, D));
  tau ~ std_normal();
  tau_int ~ normal(D*1.00, 1); // in expectation, the 'internal' weights part has a variance of D along the diagonals, before being rescaled by tau
  to_vector(theta_free) ~ std_normal();
  for (d in 1:D){
    gamma[d] ~ normal(1, 1);
  }
  L ~ lkj_corr_cholesky(1); // uniform dist
  for (n in 1:N){
    Y[jj[n],ii[n]] ~ bernoulli(Phi_approx((theta[jj[n],] * weights * bk[ii[n],]') + ak[ii[n]]));
  }
}
generated quantities{
  vector[D] r_sq; // share of var in final theta explained by lpm
  vector[D] r_old; // correlation of final theta with theta at t-1
  real r_dim; // correlation of final theta across dims
  real m_prop_compl; // mean proposal complexity
  real sd_prop_compl; // sd proposal complexity
  real pred = 0.0;
  real p_share; // share of correct votes
  matrix[J,D] theta_pred;
  matrix[J,D] theta_resid;
  row_vector[K+T] gamma_post[D];
  vector[N] log_lik;
  
  r_dim = calc_cor(theta[,1], theta[,2]);
  
  for (d in 1:D){
    theta_pred[,d] = to_vector(theta_in[,1,d])*(gamma[d]) + to_vector(theta_in[,2,d]); 
    theta_resid[,d] = (theta_pred[,d] - mean(theta_pred[,d])) - (theta_t[,d] - mean(theta_t[,d]));
    r_sq[d] = sd(theta_pred[,d])*sd(theta_pred[,d]) / (sd(theta_pred[,d])*sd(theta_pred[,d]) + sd(theta_resid[,d])*sd(theta_resid[,d]));
    
    if (T == 0){
      r_old[d] = 0.0;
    } else{
      r_old[d] = calc_cor(to_vector(theta_old[,d]), theta[theta_old_ind,d]);
    }
    if (T == 0){
      gamma_post[d] = to_row_vector((append_col(rep_vector(1,J),to_matrix(to_vector(theta_in[,1,d])))'*append_col(rep_vector(1,J),to_matrix(to_vector(theta_in[,1,d])))) \ append_col(rep_vector(1,J),to_matrix(to_vector(theta_in[,1,d])))' * theta[,d]); 
    } else {
      gamma_post[d] = to_row_vector((append_col(rep_vector(1,J),to_matrix(theta_in[,,d]))'*append_col(rep_vector(1,J),to_matrix(theta_in[,,d]))) \ append_col(rep_vector(1,J),to_matrix(theta_in[,,d]))' * theta[,d]);
    }

  }
  
  {
    vector[I] prop_compl;
    for (i in 1:I) prop_compl[i] = (pow_sum(bk[i,]', 2.00))^2 / pow_sum(bk[i,]', 4.00); // Hofmann, R. J. (1977). Indices descriptive of factor complexity. Journal of General Psychology, 96(1), 103–110. 
    m_prop_compl = mean(prop_compl);
    sd_prop_compl = sd(prop_compl);
  }

  for (n in 1:N){
    log_lik[n] = bernoulli_lpmf(Y[jj[n],ii[n]] | Phi_approx((theta[jj[n],] * weights * bk[ii[n],]') + ak[ii[n]]));
    if (Y[jj[n],ii[n]] == bernoulli_rng(Phi_approx((theta[jj[n],] * weights * bk[ii[n],]') + ak[ii[n]]))){
           pred += 1.00;
         }
  }
  p_share = pred / (I*J*1.00);
}